from os import getcwd
import os.path as ospath
from sys import argv, exit
from PyQt5.QtCore import QObject, pyqtSignal, pyqtSlot, Qt
from PyQt5.QtWidgets import (QTextEdit, QSplitter, QWidget, QPushButton, QVBoxLayout, QFileDialog, QMessageBox, QApplication)
from PyQt5.QtGui import QFont, QIcon
from msgList import MsgList
from flowlayout import FlowLayout
from random import choice as random_choice
from re import M, I
from re import split as re_split
from re import search as re_search
from torch import tensor, argsort
from torch import load as torch_load
import Model.data_loader as data_loader
from Model.config import BasicConfig
from Model.run import Runner

DEFAULT_HEAD = 'icons/user.png'
SYSTEM_HEAD = 'icons/system.png'


class TextEdit(QTextEdit, QObject):
    '''支持ctrl+return信号发射的QTextEdit'''
    entered = pyqtSignal()
    def __init__(self, parent=None):
        super(TextEdit, self).__init__(parent)
    
    def keyPressEvent(self, e):
        # print e.key() == Qt.Key_Return,e.key() == Qt.Key_Enter, e.modifiers() == Qt.ControlModifier
        if (e.key() == Qt.Key_Return) and (e.modifiers() == Qt.ControlModifier):
            self.entered.emit()# ctrl+return 输入
            self.clear()
        super(TextEdit,self).keyPressEvent(e)


'''
"QPushButton{background-color:black;color: white;border-radius: 10px;border: 2px groove gray;border-style: outset;}"
"QPushButton:hover{background-color:white; color: black;}"
"QPushButton:pressed{background-color:rgb(85, 170, 255);border-style: inset;}"
'''


class MsgInput(QWidget, QObject):
    '''自定义的内容输入控件，支持图像和文字的输入，文字输入按回车确认。'''
    textEntered = pyqtSignal(str)
    imgEntered = pyqtSignal(str)

    btnSize = 35
    teditHeight = 200

    def __init__(self, model_runner, parent=None):
        super(MsgInput, self).__init__(parent)
        self.setContentsMargins(3, 3, 3, 3)

        self.model_runner = model_runner

        self.textEdit = TextEdit()
        self.textEdit.setMaximumHeight(self.teditHeight)
        self.setMaximumHeight(self.teditHeight + self.btnSize)
        self.textEdit.setFont(QFont("Times", 20, QFont.Normal))
        self.textEdit.entered.connect(self.sendText)

        self.sendTxtBtn = QPushButton('发送')
        self.sendTxtBtn.setFont(QFont("Microsoft YaHei", 12, QFont.Bold))
        self.sendTxtBtn.setFixedHeight(self.btnSize)
        self.sendTxtBtn.clicked.connect(self.sendText)
        self.sendTxtBtn.setShortcut('enter')

        self.trainBtn = QPushButton('训练')
        self.trainBtn.setFont(QFont("Microsoft YaHei", 12, QFont.Bold))
        self.trainBtn.setFixedHeight(self.btnSize)
        # self.trainBtn.clicked.connect(self.trainModel)

        hl = FlowLayout()
        # hl.addWidget(sendImg)
        hl.addWidget(self.sendTxtBtn)
        hl.addWidget(self.trainBtn)

        vl = QVBoxLayout()
        vl.addWidget(self.textEdit)
        vl.addLayout(hl)
        self.setLayout(vl)

    def sendImage(self):  # 选择图像发送
        dialog = QFileDialog(self, '请选择图像文件...')
        dialog.setDirectory(getcwd() + '/ref')  # 设置默认路径
        dialog.setNameFilter("图片文件(*.png *.jpg *.bmp *.ico);;")  # 中间一定要用两个分号才行！
        if dialog.exec_():
            selectFileName = dialog.selectedFiles()[0]
            self.imgEntered.emit(selectFileName)
        else:  # 放弃选择
            pass

    def sendText(self):
        txt = self.textEdit.toPlainText()
        if len(txt) > 0:
            self.textEntered.emit(txt)

# class AnalysisThread(QThread):
#     analysisSignal = pyqtSignal(str)
#
#     def __init__(self, txt):
#         super().__init__()
#         self.txt = txt
#
#     def run(self):
#         self.analysisSignal.emit(self.txt)

class PyqtChatApp(QSplitter):
    """聊天界面，QSplitter用于让界面可以鼠标拖动调节"""
    curUser = {'id': None, 'name': None, 'head': DEFAULT_HEAD}
    selfHead = DEFAULT_HEAD

    def __init__(self, config):
        super(PyqtChatApp, self).__init__(Qt.Horizontal)
        print('开始初始化模型...')
        self.config = config
        self.init_model()
        self.model_runner = Runner(config)
        self.new_entities = []
        self.new_relations = []
        self.new_triples = []

        self.setWindowTitle('pyChat') # window标题
        self.setWindowIcon(QIcon('icons/chat.png')) #ICON
        self.setMinimumSize(1000, 700) # 窗口最小大小

        self.msgList = MsgList()
        self.msgInput = MsgInput(self.model_runner)
        self.msgInput.textEntered.connect(self.sendTextMsg)
        # self.msgInput.imgEntered.connect(self.sengImgMsg)
        self.msgInput.trainBtn.clicked.connect(self.train_model)
        
        rSpliter = QSplitter(Qt.Vertical, self)
        self.msgList.setParent(rSpliter)
        self.msgInput.setParent(rSpliter)

        self.show()

        hint = '请按照以下格式输入：\n 1.加入知识：如果*；进行操作*；则会* \n 2.加入知识：如果*；则进行* \n ' \
               '例如： \'加入知识：如果电网规模较大；则预设平衡点\''
        self.msgList.addTextMsg(hint, True, SYSTEM_HEAD)
        self.nonsenses = ['输入格式难倒我了', '输入格式触碰到我的知识盲区了', '我还不会这种输入格式', '换个输入呗', '注意看提示', '请按给定格式输入', hint]
        self.busy = ['正在后台忙', '正在后台处理', '等一下啦', '稍微等一下', '我在忙的', '等一下，理解万岁']
        self.state = 0

    def init_model(self):
        self.data_path = self.config.data_path
        self.result_path = self.config.result_path
        self.load_model()
        self.topk = self.config.chosen_topk
        self.entities = self.readData(ospath.join(self.data_path, 'entities.txt'))
        entity_ids, entity_lens = [], []
        for entity in self.entities:
            entity_id, e_len = data_loader.TrainDataset.sequence(data_loader.TrainDataset.tokenizer, entity, self.config.max_seq_len)
            entity_ids.append(entity_id)
            entity_lens.append(e_len)
        self.relations = self.readData(ospath.join(self.data_path, 'relations.txt'))
        relation_ids, relation_lens = [], []
        for relation in self.relations:
            relation_id, r_len = data_loader.TrainDataset.sequence(data_loader.TrainDataset.tokenizer, relation, self.config.max_seq_len)
            relation_ids.append(relation_id)
            relation_lens.append(r_len)
        self.entity_ids, self.entity_lens = tensor(entity_ids).long().to(self.config.device), entity_lens
        self.relation_ids, self.relation_lens = tensor(relation_ids).long().to(self.config.device), relation_lens

    @pyqtSlot(str)
    def sendTextMsg(self, txt):
        self.msgList.addTextMsg(txt, False)
        # self.analysisThread = AnalysisThread(txt)
        # self.analysisThread.analysisSignal.connect(self.process)
        # self.analysisThread.start()
        self.msgInput.textEdit.clear()
        # print('子线程开始')
        self.process(txt)
        # self.analysis(txt)
        # self.msgList.addTextMsg(self.msgList.r_text_history[-1], True, SYSTEM_HEAD)

    def process(self, txt):
        if txt.lower() in ['结束', 'end']:
            if self.state == 0:
                self.msgList.addTextMsg('还没开始啊', True, SYSTEM_HEAD)
            else:
                self.msgList.addTextMsg('收到，这一轮结束，以上内容不会写入数据库', True, SYSTEM_HEAD)
            return
        if self.state == 1:
            self.msgList.addTextMsg(random_choice(self.busy), True, SYSTEM_HEAD)
        elif self.state == 0:
            if txt[:2] not in ['加入', '插入', '增加']:
                self.msgList.addTextMsg(random_choice(self.nonsenses), True, SYSTEM_HEAD)
                return
            self.state = 1
            self.msgInput.sendTxtBtn.setDisabled(True)
            QApplication.processEvents()
            try:
                self.analysis(txt)
            except:
                self.msgList.addTextMsg('请按照格式输入！', True, SYSTEM_HEAD)
                self.state = 0
            self.msgInput.sendTxtBtn.setDisabled(False)
            QApplication.processEvents()
        elif self.state == 2:
            if txt.lower() in ['no', '有错', '不对', '错误']:
                sz = '请给出错误部分您的定义'
                self.msgList.addTextMsg(sz, True, SYSTEM_HEAD)
                sz = '（不同部分内容用回车符隔开，以head和relation自定义为例，格式示例：\n' \
                     '\'head: %用户定义实体%\nrelation: %用户定义关系%\'）'
                self.msgList.addTextMsg(sz, True, SYSTEM_HEAD)

                self.state = 3
            else:
                try:
                    span = re_search(r'[0-9]+ [0-9]+ [0-9]+', txt, M|I).span()
                    if span is None:
                        self.msgList.addTextMsg(random_choice(self.nonsenses), True, SYSTEM_HEAD)
                    else:
                        elems = txt[span[0]: span[1]].strip().split(' ')
                        elems = [int(elem) for elem in elems]
                        try:
                            self.head = self.heads[elems[0] - 1]
                            self.relation = self.properties[elems[1] - 1]
                            self.tail = self.tails[elems[2] - 1]
                            self.affirm(self.head, self.relation, self.tail)
                            self.state = 4
                        except:
                            self.msgList.addTextMsg('请给出正确范围内的序号', True, SYSTEM_HEAD)
                except:
                    self.msgList.addTextMsg('请按格式输入', True, SYSTEM_HEAD)
        elif self.state == 3:
            phases = txt.strip().split('\n')
            phases = [phase for phase in phases if phase]
            if phases:
                part_value = {}
                try:
                    for phase in phases:
                        elems = re_split(':|：', phase.strip())
                        part_value[elems[0]] = elems[1]
                    for part in part_value:
                        if part.lower() in ['head', 'h', '头实体', '头结点']:
                            self.heads = [part_value[part]]
                            # self.new_entities.append([self.head_str, part_value[part]])
                        elif part.lower() in ['relation', 'r', '关系', '谓语']:
                            self.properties = [part_value[part]]
                            # self.new_relations.append([self.relation_str, part_value[part]])
                        elif part.lower() in ['tail', 't', '尾节点', '尾实体']:
                            self.tails = [part_value[part]]
                            # self.new_entities.append([self.tail_str, part_value[part]])
                        else:
                            self.msgList.addTextMsg('请按格式输入', True, SYSTEM_HEAD)
                            return
                    if self.heads is None or self.properties is None or self.tails is None:
                        self.queryForDefinition(self.heads, self.properties, self.tails)
                    else:
                        self.repetition(self.heads, self.properties, self.tails)
                        self.state = 2
                except:
                    self.msgList.addTextMsg(random_choice(self.nonsenses), True, SYSTEM_HEAD)
                    self.msgList.addTextMsg('格式不对，请重新输入', True, SYSTEM_HEAD)
                    return
            else:
                self.msgList.addTextMsg('请按格式输入', True, SYSTEM_HEAD)
        elif self.state == 4:
            if txt.lower() in ['yes', '对', '对的', '没错', '正确', '可以']:
                self.new_entities.append([self.head_str, self.head])
                self.new_entities.append([self.tail_str, self.tail])
                self.new_relations.append([self.relation_str, self.relation])
                self.new_triples.append([self.head, self.relation, self.tail])
                self.writeKnowledge()
                self.msgList.addTextMsg('好的，我记住了，这一轮结束', True, SYSTEM_HEAD)
                self.state = 0
            elif txt.lower() in ['no', '不是', '错', '不对']:
                self.msgList.addTextMsg('好吧，那重新选择看看有没有合适的', True, SYSTEM_HEAD)
                self.repetition(self.heads, self.properties, self.tails)
                self.state = 2
            else:
                self.msgList.addTextMsg(random_choice(self.nonsenses), True, SYSTEM_HEAD)
                self.msgList.addTextMsg('格式不对，请重新输入', True, SYSTEM_HEAD)

    def writeKnowledge(self):
        '''
        TODO
        '''
        self.writeNewTriples()
        print('写入完成...')
        self.new_triples.clear()

        if len(self.new_entities) + len(self.new_relations) >= self.config.new_triples_threshold:
            self.train()

    def train_model(self):
        confirm = QMessageBox.information(self, "提示", "确认启动重训练吗？", QMessageBox.Yes, QMessageBox.No)
        if confirm == QMessageBox.Yes:
            self.train()

            print('训练完毕...')

    def repetition(self, heads, relations, tails):
        txt = '识别出三元组 (head, relation, tail)：\nhead：\n'
        for idx, head in enumerate(heads):
            txt += '{}. {}\n'.format(idx + 1, head)
        txt += 'relation：\n'
        for idx, relation in enumerate(relations):
            txt += '{}. {}\n'.format(idx + 1, relation)
        txt += 'tail：\n'
        for idx, tail in enumerate(tails):
            txt += '{}. {}\n'.format(idx + 1, tail)
        txt += '请选择正确的head、relation、tail 的序号（格式示例：\'1 2 1\'）'
        self.msgList.addTextMsg(txt, True, SYSTEM_HEAD)
        self.msgList.addTextMsg('注意：如果存在识别错误的部分，请输入 no', True, SYSTEM_HEAD)

    def affirm(self, head, relation, tail):
        txt = '确认是知识（{}；{}；{}）\n吗？请输入 yes/no'.format(head, relation, tail)
        self.msgList.addTextMsg(txt, True, SYSTEM_HEAD)

    def queryForDefinition(self, heads, relations, tails):
        is_none = ''
        if heads is None:
            is_none += 'head、'
        if relations is None:
            is_none += 'relation、'
        if tails is None:
            is_none += 'tail'
        is_none = is_none[:-1]

        txt = '{} 部分未识别出，可能并未在数据库中定义，\n' \
              '需要您自行定义，请给出以上内容的定义\n' \
              '（不同部分内容用回车符隔开，\n' \
              '以head和relation自定义为例，格式示例：\n' \
              '\'head: %用户定义实体%\n' \
              'relation: %用户定义关系%\'）：'.format(is_none)
        self.msgList.addTextMsg(txt, True, SYSTEM_HEAD)

    def analysis(self, text):
        text = [elem for elem in re_split('：|:', text.strip()) if elem]
        assert len(text) == 2
        text = text[1]
        phases = re_split('；|;', text)
        self.head_str = phases[0].replace('如果', '').replace('当', '')
        head_ids, head_len = data_loader.TrainDataset.sequence(data_loader.TrainDataset.tokenizer, self.head_str, self.config.max_seq_len)
        head_ids = tensor(head_ids).long().unsqueeze(0).repeat(len(self.entity_lens), 1).to(self.config.device)
        head_len = [head_len for _ in range(len(self.entity_lens))]
        predicts = self.model.forward(head_ids, self.entity_ids, head_len, self.entity_lens)[:, 1]
        chosen_head_indices = self.select(predicts, self.topk)
        self.heads = None if len(chosen_head_indices) == 0 else [self.entities[indice] for indice in chosen_head_indices]

        if len(phases) == 2:
            self.properties = ['操作']
        else:
            self.relation_str = phases[1].replace('进行', '').replace('操作', '')
            relation_ids, relation_len = data_loader.TrainDataset.sequence(data_loader.TrainDataset.tokenizer, self.relation_str, self.config.max_seq_len)
            relation_ids = tensor(relation_ids).long().unsqueeze(0).repeat(len(self.relation_lens), 1).to(self.config.device)
            relation_len = [relation_len for _ in range(len(self.relation_lens))]
            predicts = self.model.forward(relation_ids, self.relation_ids, relation_len, self.relation_lens)[:, 1]
            chosen_relation_indices = self.select(predicts, self.topk)
            self.properties = None if len(chosen_relation_indices) == 0 else [self.relations[indice] for indice in chosen_relation_indices]

        self.tail_str = phases[-1].replace('则', '').replace('会', '')
        tail_ids, tail_len = data_loader.TrainDataset.sequence(data_loader.TrainDataset.tokenizer, self.tail_str,
                                                               self.config.max_seq_len)
        if len(phases) == 3:
            tail_ids = tensor(tail_ids).long().unsqueeze(0).repeat(len(self.entity_lens), 1).to(self.config.device)
            tail_len = [tail_len for _ in range(len(self.entity_lens))]
            predicts = self.model.forward(tail_ids, self.entity_ids, tail_len, self.entity_lens)[:, 1]
            chosen_tail_indices = self.select(predicts, self.topk)
            self.tails = None if len(chosen_tail_indices) == 0 else [self.entities[indice] for indice in
                                                                     chosen_tail_indices]
        else:
            tail_ids = tensor(tail_ids).long().unsqueeze(0).repeat(len(self.relation_lens), 1).to(
                self.config.device)
            tail_len = [tail_len for _ in range(len(self.relation_lens))]
            predicts = self.model.forward(tail_ids, self.relation_ids, tail_len, self.relation_lens)[:, 1]
            chosen_tail_indices = self.select(predicts, self.topk)
            self.tails = None if len(chosen_tail_indices) == 0 else [self.relations[indice] for indice in
                                                                     chosen_tail_indices]

        if self.heads is None or self.properties is None or self.tails is None:
            self.msgList.addTextMsg('能力不足，有部分没识别出来', True, SYSTEM_HEAD)
            self.queryForDefinition(self.heads, self.properties, self.tails)
            self.state = 3
        else:
            self.repetition(self.heads, self.properties, self.tails)
            self.state = 2

    def select(self, values, topk):
        # values_i = torch.argsort(values, descending=True)
        flag = (values > 0.5).nonzero().squeeze(dim=1)
        satisfy_values = values[flag]
        values_i = argsort(satisfy_values, descending=True)[:topk]
        return flag[values_i]

    def load_model(self):
        if ospath.exists(ospath.join(self.result_path, "model")):
            self.model = torch_load(ospath.join(self.result_path, "model"))
            # self.optimizer = torch.load(os.path.join(self.result_path, "optimizer"))
        else:
            raise ValueError('Model not exist...')

    def writeNewTriples(self):
        with open(ospath.join(self.data_path, 'triples_str.txt'), 'a+', encoding="utf-8") as fout:
            for triple in self.new_triples:
                fout.write('\t'.join(triple) + '\n')

    def writeNewEntitiesAndRelations(self):
        with open(ospath.join(self.data_path, 'train_entities.txt'), 'a+', encoding="utf-8") as fout:
            for entity_pair in self.new_entities:
                fout.write('\t'.join(entity_pair) + '\n')
        with open(ospath.join(self.data_path, 'train_relations.txt'), 'a+', encoding='utf-8') as fout:
            for relation_pair in self.new_relations:
                fout.write('\t'.join(relation_pair) + '\n')

        with open(ospath.join(self.data_path, 'entities.txt'), 'a+', encoding="utf-8") as fout:
            for entity in self.new_entities:
                if entity[1] not in self.entities:
                    fout.write(entity[1] + '\n')
        with open(ospath.join(self.data_path, 'relations.txt'), 'a+', encoding="utf-8") as fout:
            for relation in self.new_relations:
                if relation[1] not in self.relations:
                    fout.write(relation[1] + '\n')

    def readData(self, path):
        contents = []
        with open(path, 'r', encoding="utf-8") as fin:
            for line in fin:
                contents.append(line.strip())
        return contents

    def train(self):
        self.writeNewEntitiesAndRelations()
        print('清除缓存...')
        self.new_entities.clear()
        self.new_relations.clear()
        print("开始训练...")

        self.msgInput.sendTxtBtn.setDisabled(True)
        self.msgInput.trainBtn.setDisabled(True)
        self.model_runner.run()
        self.msgInput.sendTxtBtn.setDisabled(False)
        self.msgInput.trainBtn.setDisabled(False)

        print('重新初始化模型')
        self.init_model()





if __name__=='__main__':
    # 加入：如果电网有直流系统，进行减少分区内发电机的有功功率操作，则处于平衡状态
    # head:电网有直流系统
    # relation:减少分区内发电机的有功功率
    # tail:处于平衡状态

    data_path = ospath.join(".", "data", "PowerGrid")
    icon_path = ospath.join(".", 'icons')
    record_path = ospath.join(".", "log", "PowerGrid")
    result_path = ospath.join(".", "result", "PowerGrid")
    use_cuda = False
    do_train, do_valid = True, False
    use_checkpoint = True

    hidden_dim = 32
    num_perspective = 16
    chosen_topk = 3

    lr = 0.0005
    cpu_num = 1
    train_batch_size = 20
    test_batch_size = 4
    epoch_num = 1
    train_max_steps = 10
    warm_up_step = 100
    log_step = 10
    valid_step = 100
    max_seq_len = 30
    negative_sample_size = 1
    checkpoint_step = 100
    embedding_dim = 768
    output_dim = 16

    new_triples_threshold = 2      # 当新增三元组（实体或者关系之前未遇到过）超过这个值后系统自动开始训练

    config = BasicConfig(data_path, record_path, result_path,
                         use_cuda, do_train, do_valid, use_checkpoint, lr, cpu_num,
                         train_batch_size, test_batch_size,
                         epoch_num, train_max_steps, warm_up_step, log_step,
                         valid_step, max_seq_len, negative_sample_size,
                         checkpoint_step,
                         embedding_dim=embedding_dim, chosen_topk=chosen_topk)
    config.new_triples_threshold = new_triples_threshold
    config.icon_path = icon_path

    app = QApplication(argv)
    pchat = PyqtChatApp(config)
    exit(app.exec_())

    # 插入知识：如果初始状态；进行操作退出并联电容器；则电网无直流系统
    # 加入知识：如果发电机有功功率达到下限；进行操作检查各分区的无功发电、无功负荷是否平衡；则发电机有功功率在正常范围
    # 加入知识：如果计算不收敛；则对电网进行简化等值，判断电网是否有新能源发电
    # 加入知识：如果电网无新能源发电；进行操作调回平衡状态；则平衡状态